import torch
from torch import nn
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# 读取训练集数据
train_dataframe = pd.read_csv('一刀切型散点' + '\\' + 'train.csv', header=None)
# 将dataFrame类型的训练数据train_dataframe转换为numpy数组类型。
train_data = train_dataframe.values
# x为数组的第一行，y为数组的第二行，label为数组的第三行
x = train_data[0]
y = train_data[1]
# 将输入数据x和y组成一个数组，并将其转换为numpy数组类型，数组内元素类型转换为单精度浮点数类型
inputs = np.array([x, y]).astype('float32')
label = np.array(train_data[2]).astype('float32')
# 可视化散点图
plt.scatter(x, y, c=label)
plt.show()

# 定义神经网络模型
class Net(nn.Module):
    # 初始化模型
    def __init__(self):
        super(Net, self).__init__()
        # 输入层有x,y两个特征参数，输出层有labels一个神经元，设隐藏层有5个神经元
        # 用nn.Linear连接输入层和隐藏层，隐藏层和输出层
        self.fc_layer1 = nn.Linear(2, 10)
        self.fc_layer2 = nn.Linear(10, 1)

    # 前馈函数
    def forward(self, x):
        # 将第一个全连接层的输出结果通过ReLU激活函数进行非线性映射
        x = torch.relu(self.fc_layer1(x.T))
        # 将经过ReLU激活后的结果再通过第二个全连接层进行线性变换
        # 并将输出结果通过Sigmoid函数进行压缩，得到一个在0-1之间的概率值，表示正样本的概率。
        x = torch.sigmoid(self.fc_layer2(x))
        return x

# 创建模型
net = Net()
# 损失函数为均方误差，较符合该数据集
criterion = nn.MSELoss()
# 优化器
optimizer = torch.optim.SGD(net.parameters(), lr=0.10)

# 训练
for iteration in range(1000):
    inputs = np.array([x, y]).astype('float32')
    labels = np.array(label).astype('float32')
    # 将inputs转换为torch张量
    inputs = torch.from_numpy(inputs)
    # torch.from_numpy()函数用来将numpy数组转换为torch张量，并且这两者共享内存
    # 将labels转换为torch张量
    labels = torch.from_numpy(labels)
    # 用定义好的神经网络模型 net 对输入数据 inputs 进行前向传播计算，得到模型的输出结果 outputs
    outputs = net(inputs)
    # criterion会返回一个标量值给loss，表示模型在当前批次上的平均损失值
    loss = criterion(outputs.squeeze(), labels)
    # loss.backward() 是 PyTorch 中用于计算模型参数梯度的函数，通过反向传播机制将当前损失值向输出端进行传播
    # 并计算出所有需要计算的参数的梯度值，以便进行优化和参数更新。
    loss.backward()
    # 使用optimizer.step()更新网络参数
    optimizer.step()

    if iteration % 100 == 0:
        # 每迭代100次输出一次损失函数
        print("iteration:{}  loss:{}".format(iteration, loss.item()))

# 测试
test_dataframe = pd.read_csv('一刀切型散点' + '\\' + 'test.csv', header=None)
test_data = test_dataframe.values
x = test_data[0]
y = test_data[1]
inputs = np.array([x, y]).astype('float32')
labels = np.array(test_data[2]).astype('float32')
inputs = torch.from_numpy(inputs)
labels = torch.from_numpy(labels)
# 将测试集的样本数据传入模型中进行测试，得到输出结果
outputs = net(inputs)
_, predicted = torch.max(outputs, 1)
predicted = (outputs > 0.5).float().squeeze()
# squeeze()会让张量去掉维数为1的维度
total = predicted.size(0)
correct = (predicted == labels).sum().item()
accuracy = correct / total
print('Accuracy: {}%'.format(accuracy * 100))
